package com.mf.utils.tree;

import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import org.apache.commons.lang3.ObjectUtils;

import java.util.*;
import java.util.stream.Collectors;

/**
 * @Author: lmf
 * @Create: 2025/4/5 10:10
 * @module 列表根据pid字段构建树形结构
 */
public class UtilTree {
    /**
     * 顶级pid
     */
    public static final Integer ROOT_PID = 0;

    /**
     * 构建树结构
     */
    public static <T extends TreeNode<T, K>, K> List<T> buildTree(List<T> treeNode) {
        HashMap<K, T> nodeMap = new HashMap<>();
        for (T t : treeNode) {
            nodeMap.put(t.getId(), t);
        }

        // 找出所有的父节点 pid = 0 父节点 ->  0 || “0”
        List<T> rootNodes = new ArrayList<>();
        for (T node : treeNode) {
            if (node.getPid() instanceof Integer || node.getPid() instanceof String) {
                if (node.getPid().equals(ROOT_PID)) {
                    rootNodes.add(node);
                }
            }
        }

        // 递归构建树结构
        for (T rootNode : rootNodes) {
            buildTree(rootNode, nodeMap);
        }
        // 对构建出来的树结构按照sort字段进行排序
        rootNodes = sortTree(rootNodes);
        return rootNodes;
    }

    /**
     * @param list<T> 所有数据的树结构
     * @param pageList   需要显示的节点id
     */
    public static <T extends TreeNode<T, K>, K> Page<T> filterTree(List<T> list, Page<T> pageList) {
        List<T> ts = filterTree(list, pageList.getRecords());
        pageList.setRecords(ts);
        // 重新计算页数
        pageList.setTotal(ts.size());
        // 重新计算总条数
        pageList.setPages((pageList.getTotal() + pageList.getSize() - 1) / pageList.getSize());
        return pageList;
    }


    /**
     * @param list       所有数据的树结构
     * @param targetList 需要显示的节点集合
     */
    public static <T extends TreeNode<T, K>, K> List<T> filterTree(List<T> list, List<T> targetList) {
        Map<K, T> targetMap = targetList.stream().collect(Collectors.toMap(TreeNode::getId, t -> t));
        // 替换 list 中与 targetList 同 ID 的节点 避免如果targetList是连表等查询有些字段丢失
        List<T> updateList = list.stream().map(node -> {
            if (targetMap.containsKey(node.getId())) {
                return targetMap.get(node.getId());
            }
            return node;
        }).toList();

        Set<K> targetIds = targetList.stream().map(TreeNode::getId).collect(Collectors.toSet());
        // 构建 nodeMap 用于快速查找父节点
        Map<K, T> nodeMap = new HashMap<>();
        for (T t : updateList) {
            nodeMap.put(t.getId(), t);
        }

        // 构建 parentMap 用于查找子节点
        Map<K, List<T>> parentMap = updateList.stream()
                .collect(Collectors.groupingBy(TreeNode::getPid));

        // 收集所有需要保留的 ID（包括目标、子孙、祖先）
        Set<K> allTargetIds = new HashSet<>();

        for (K id : targetIds) {
            allTargetIds.add(id); // 目标节点
            allTargetIds.addAll(collectDescendantIds(id, parentMap)); // 子孙
            allTargetIds.addAll(collectAncestorIds(id, nodeMap, null)); // 祖先
        }

        // 过滤出需要保留的节点
        List<T> result = updateList.stream()
                .filter(t -> allTargetIds.contains(t.getId()))
                .collect(Collectors.toList());

        // 构建并返回树结构
        return buildTree(result);
    }

    /**
     * 递归收集目标节点的所有祖先节点的 ID
     */
    private static <T extends TreeNode<T, K>, K> Set<K> collectAncestorIds(
            K id,
            Map<K, T> nodeMap,
            Set<K> visited) {

        if (visited == null) {
            visited = new HashSet<>();
        }

        if (id == null || visited.contains(id)) {
            return visited;
        }

        visited.add(id);

        T node = nodeMap.get(id);
        if (node != null) {
            K pid = node.getPid();
            if (pid != null && !pid.equals(ROOT_PID)) {
                collectAncestorIds(pid, nodeMap, visited);
            }
        }

        return visited;
    }

    /**
     * 收集指定节点下的所有子孙节点 ID（递归 + 广度优先遍历）
     *
     * @param id       当前节点 ID（起始节点）
     * @param parentMap 父节点映射表（key=父ID，value=子节点列表）
     * @return 所有子孙节点 ID 集合
     */
    private static <T extends TreeNode<T, K>, K> Set<K> collectDescendantIds(K id, Map<K, List<T>> parentMap) {
        // 存储所有子孙节点的 ID（避免重复）
        Set<K> descendants = new HashSet<>();

        // 使用队列实现广度优先遍历（BFS）
        Queue<K> queue = new LinkedList<>();
        queue.add(id); // 起始节点加入队列

        while (!queue.isEmpty()) {
            K currentId = queue.poll(); // 取出当前要处理的节点 ID

            // 获取该节点的所有直接子节点
            List<T> children = parentMap.getOrDefault(currentId, Collections.emptyList());

            for (T child : children) {
                K childId = child.getId(); // 子节点 ID

                // 如果是第一次遇到这个子节点（未被添加过）
                if (descendants.add(childId)) {
                    // 就把它加入队列，继续查找其子节点
                    queue.add(childId);
                }
            }
        }

        return descendants; // 返回所有子孙节点 ID
    }


    /**
     * 将树结构平铺为一维数组 （主要为了让部分先按照树的结构进行排序，这样在平铺出来的就是排序后的结构）
     */
    public static <T extends TreeNode<T, K>, K> List<T> flattenTree(List<T> treeNode) {
        List<T> flattenedTree = new ArrayList<>();
        for (T node : treeNode) {
            flattenedTree.add(node);
            if (ObjectUtils.isNotEmpty(node.getChildren())) {
                flattenedTree.addAll(flattenTree(node.getChildren()));
            }
        }
        return flattenedTree;
    }

    /**
     * 递归构建树结构
     */
    private static <T extends TreeNode<T, K>, K> void buildTree(T rootNode, HashMap<K, T> nodeMap) {
        for (T node : nodeMap.values()) {
            if (node.getPid().equals(rootNode.getId())) {
                rootNode.addChildren(node);
                buildTree(node, nodeMap);
            }
        }
    }

    /**
     * 对构建出来的树结构按照sort字段进行排序
     */
    private static final Comparator<?> ID_COMPARATOR = (o1, o2) -> {
        if (o1 instanceof String && o2 instanceof String) {
            return ((String) o1).compareTo((String) o2);
        } else if (o1 instanceof Number && o2 instanceof Number) {
            return Double.compare(((Number) o1).doubleValue(), ((Number) o2).doubleValue());
        } else {
            return o1.toString().compareTo(o2.toString());
        }
    };

    @SuppressWarnings("unchecked")
    private static <T extends TreeNode<T, K>, K> List<T> sortTree(List<T> treeNode) {
        Comparator<T> comparator = Comparator.<T, Integer>comparing(TreeNode::getSort).thenComparing(TreeNode::getId,
                (Comparator<K>) ID_COMPARATOR);

        treeNode = treeNode.stream().sorted(comparator.reversed()).toList();

        for (T node : treeNode) {
            if (ObjectUtils.isNotEmpty(node.getChildren())) {
                node.setChildren(sortTree(node.getChildren()));
            }
        }

        return treeNode;
    }


}
